import os
import requests
from volcenginesdkarkruntime import Ark

VOLC_EMBEDDINGS_API_URL = os.getenv("VOLC_EMBEDDINGS_API_URL")
VOLC_EMBEDDINGS_VISION_API_URL = os.getenv("VOLC_EMBEDDINGS_VISION_API_URL")
VOLC_API_KEY = os.getenv("VOLC_API_KEY")


def get_doc_embedding(doc_content):
    client = Ark(api_key=VOLC_API_KEY)

    response = client.embeddings.create(
        model="doubao-embedding-text-240715",
        input=[doc_content],
        encoding_format="float",
    )
    res = response.model_dump()
    return res["data"][0]["embedding"]


def get_doc_embedding2(doc_content):
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {VOLC_API_KEY}",
    }
    payload = {
        "encoding_format": "float",
        "input": [doc_content],
        "model": "doubao-embedding-text-240715",
    }
    response = requests.post(VOLC_EMBEDDINGS_API_URL, json=payload, headers=headers)
    if response.status_code == 200:
        data = response.json()
        return data["data"][0]["embedding"]
    else:
        raise Exception(f"Embedding API error:{response.text}")


def get_image_embedding(image_url):
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {VOLC_API_KEY}",
    }
    payload = {
        "encoding_format": "float",
        "input": [{"type": "image_url", "image_url": {"url": image_url}}],
        "model": "doubao-embedding-vision-250615",
    }
    response = requests.post(
        VOLC_EMBEDDINGS_VISION_API_URL, json=payload, headers=headers
    )
    if response.status_code == 200:
        data = response.json()
        return data["data"]["embedding"]
    else:
        raise Exception(f"Embedding API error:{response.text}")
